#!/usr/bin/python
# -*-coding:utf-8-*-
import os
import pandas as pd

from base.GPModelSave import load_gpmodel
from base.GPModelSave import model_predict
from statsmodels.regression.linear_model import OLS

from base.KeyTransform import stock_code_map
from base.KeyTransform import trade_date_map

barra_data_dir = './jq_barra_data'

def get_factor_by_model(model,
                        model_name,
                        start_date=None,
                        end_date=None,
                        raw_data=None,
                        keys=None,
                        neu_keys=None,
                        by_name=True,
                        key_by_path=False,
                        ts_base=5,
                        key_reverse=True):
    # model = load_gpmodel(os.path.join(model_dir, model_filname+'.pkl'))

    Xnames = [k for k in model if isinstance(k, str)]

    # TODO - 去掉重复的X
    Xnames = list(set(Xnames))

    used_raw_data = raw_data[Xnames]

    factor_data = model_predict(X=used_raw_data,
                                program=[model, ts_base],
                                keys=keys,
                                neu_keys=neu_keys,
                                by_name=by_name,
                                key_by_path=key_by_path)

    factor_data = pd.DataFrame(factor_data, index=keys)
    factor_data.columns = [model_name]

    factor_data = factor_data.reset_index()

    if key_reverse:
        factor_data = stock_code_map(factor_data, key='stock_code', reverse=True)
        factor_data = trade_date_map(factor_data, key='date', reverse=True)

    if start_date is None and end_date is not None:
        factor_data = factor_data[factor_data['date'] <= pd.to_datetime(end_date)]
    elif start_date is not None and end_date is None:
        factor_data = factor_data[factor_data['date'] >= pd.to_datetime(start_date)]
    elif start_date is not None and end_date is not None:
        factor_data = factor_data[(factor_data['date'] >= pd.to_datetime(start_date)) &
                                  (factor_data['date'] <= pd.to_datetime(end_date))]

    return factor_data.reset_index(drop=True)


class FactorBarraStrippingForTest(object):
    def __init__(self):
        self.jq_barra_p1, \
        self.jq_barra_p2 = self._get_barra_data()

        self.p1_barra_columns = self.jq_barra_p1.columns
        self.p2_barra_columns = self.jq_barra_p2.columns

        self.p1_end_date = '2019-11-30'


    def _get_barra_data(self):
        new_processed_jq_barra_data_v0 = pd.read_hdf(os.path.join(barra_data_dir, 'new_processed_jq_barra_data_v0.h5'))
        new_processed_jq_barra_data_v1 = pd.read_hdf(os.path.join(barra_data_dir, 'new_processed_jq_barra_data_v1.h5'))

        new_processed_jq_barra_data_v0 = new_processed_jq_barra_data_v0.set_index(['date', 'symbol'])
        new_processed_jq_barra_data_v1 = new_processed_jq_barra_data_v1.set_index(['date', 'symbol'])

        return new_processed_jq_barra_data_v0, new_processed_jq_barra_data_v1

    def factor_barra_stripping(self, factor_data, factor_name):
        ## TODO - 分段回归
        # 读取因子数据
        # p1_end_date = '2019-11-30'
        factor_data_p1 = factor_data[factor_data['date'] <= self.p1_end_date]
        factor_data_p2 = factor_data[factor_data['date'] > self.p1_end_date]

        factor_data_p1 = factor_data_p1.set_index(['date', 'stock_code'])
        factor_data_p2 = factor_data_p2.set_index(['date', 'stock_code'])

        # P1
        factor_data_p1 = pd.concat([factor_data_p1, self.jq_barra_p1], axis=1, join_axes=[factor_data_p1.index])
        factor_data_p1 = factor_data_p1.dropna()

        model = OLS(factor_data_p1[factor_name], factor_data_p1[self.p1_barra_columns], hasconst=True)

        res = model.fit(params_only=True)

        # P1剥离
        factor_data_p1[factor_name] = \
            factor_data_p1[factor_name] - factor_data_p1[self.p1_barra_columns].values @ res.params.values

        # P2
        factor_data_p2 = pd.concat([factor_data_p2, self.jq_barra_p2], axis=1, join_axes=[factor_data_p2.index])
        factor_data_p2 = factor_data_p2.dropna()

        model = OLS(factor_data_p2[factor_name], factor_data_p2[self.p2_barra_columns], hasconst=True)

        res = model.fit(params_only=True)

        # P2剥离
        factor_data_p2[factor_name] = \
            factor_data_p2[factor_name] - factor_data_p2[self.p2_barra_columns].values @ res.params.values

        # concat
        factor_data = pd.concat([factor_data_p1[[factor_name]],
                                 factor_data_p2[[factor_name]]], axis=0)

        factor_data = factor_data.reset_index()

        del factor_data_p1
        del factor_data_p2

        print(factor_name, 'barra stripped!\n')
        return factor_data


factor_barra_strip_api = FactorBarraStrippingForTest()

def get_gp_factor(raw_data,
                  keys,
                  neu_keys,
                  model,
                  factor_name,
                  by_name=True,
                  key_by_path=False,
                  ts_base=5,
                  key_reverse=True):
    raw_data = raw_data.copy()

    # TODO - 传入模型，计算因子
    factor_data = get_factor_by_model(model=model,
                                      model_name=factor_name,
                                      start_date=None,
                                      end_date=None,
                                      raw_data=raw_data,
                                      keys=keys,
                                      neu_keys=neu_keys,
                                      by_name=by_name,
                                      key_by_path=key_by_path,
                                      ts_base=ts_base,
                                      key_reverse=key_reverse)
    # print('get factor data done!')

    # TODO - 去掉nan的因子值
    factor_data = factor_data.dropna()

    # TODO - 进行barra剥离
    factor_data_barra = factor_barra_strip_api.factor_barra_stripping(factor_data=factor_data, factor_name=factor_name)
    # print('barra stripping done!')

    return factor_data, factor_data_barra



